package com.ruoyi.teaching.flink;

import com.ruoyi.common.core.utils.SpringUtils;
import com.ruoyi.teaching.domain.Interview;
import com.ruoyi.teaching.domain.YaHighInterview;
import lombok.extern.slf4j.Slf4j;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.functions.sink.RichSinkFunction;
import org.springframework.data.elasticsearch.client.elc.ElasticsearchTemplate;

import java.util.ArrayList;
import java.util.List;

import static com.ruoyi.teaching.flink.MySQLSource.BATCH_SIZE;

/**
 * @author 15332
 * @version 1.0
 * @date 2025/9/10 14:08
 * @description flink 终点
 */
@Slf4j
public class BatchESSink extends RichSinkFunction<YaHighInterview> {

  private int totalCount = 0;

  private transient ElasticsearchTemplate elasticsearchRestTemplate;

  private transient List<Interview> buffer;

  @Override
  public void open(Configuration parameters) {
    // 初始化缓冲区和Elasticsearch模板
    buffer = new ArrayList<>(BATCH_SIZE);
    elasticsearchRestTemplate = SpringUtils.getBean(ElasticsearchTemplate.class);
  }

  @Override
  public void invoke(YaHighInterview record, Context context) {
    // 将YaHighInterview对象转换为Interview对象
    Interview interview = new Interview();
    interview.setId(record.getId());
    interview.setQuestion(record.getQuestion());
    interview.setAnswer(record.getAnswer());
    interview.setKnowledgePoint(record.getKnowledgePoint());
    interview.setImportance(record.getImportance());
    interview.setStage(record.getStage());
    buffer.add(interview);
    // 当缓冲区达到批量大小时，写入Elasticsearch
    if (buffer.size() >= BATCH_SIZE) {
      writeBufferToES();
    }
  }

  private void writeBufferToES() {
    totalCount += buffer.size();
    // 将缓冲区中的数据写入Elasticsearch
    log.debug("写入新批次数据到ES，批次大小：{}", buffer.size());
    elasticsearchRestTemplate.save(buffer);
    log.debug("已写入总记录数：{}", totalCount);
    buffer.clear();
  }

  @Override
  public void close() throws Exception {
    // 在关闭前，确保缓冲区中的剩余数据也被写入Elasticsearch
    if (!buffer.isEmpty()) {
      writeBufferToES();
    }
    log.debug("全量同步到ES结束");
  }
}
